#!/usr/bin/env python
# coding=utf-8
"""
Virtual Disk Api for VMWare
"""

from os.path import join
from vmware import VM, get_vcenter_prop
import robot.libraries.BuiltIn
try:
    from pyVmomi import vim
except ImportError:
    robot.libraries.BuiltIn.BuiltIn().log('Missing pyvmomi', level='WARN')

MIN_SIZE_STEP_GB = 512
MAX_SCSI_NO = 3
MIN_SCSI_NO = 0
MAX_DEVICE_NO = 15
MIN_DEVICE_NO = 0


class VirtualDisk(VM):
    """
    Virtual Disk
    """
    def __init__(self, uuid, disk_size_gb=1 * MIN_SIZE_STEP_GB,
                 disk_provision_type='thin',
                 disk_mode='independent_persistent'):
        """
        Initialization function for VirtualDisk
        Args:
            va_lib: VALibrary object
            vm_obj: Virtual Machine Object with a certain name.
            disk_size_gb: size of the disk to be added
            disk_provision_type: Disk provision type when adding disk
                                 for VMWare.
            disk_mode: Disk mode when adding disk for VMWare
        """
        vcenter, user, password = get_vcenter_prop()
        super(VirtualDisk, self).__init__(vcenter, user, password)
        self.vm_obj = self._get_vm_by_uuid(uuid)
        self.disk_size_gb = disk_size_gb
        self.disk_provision_type = disk_provision_type
        self.disk_mode = disk_mode

    def add_virtual_disk(self, scsi_num_tuple):
        """
        Add a virtual disk to VM
        Args:
            scsi_num_tuple: scsi number which tuple format.
                            eg. (2,12) means (2:12)

        Returns: task

        """
        assert is_valid_device(scsi_num_tuple)
        # if self.disk_size_gb % MIN_SIZE_STEP_GB != 0:
        #     raise RuntimeError('Virtual disk size must be a multiple of 512 GB'
        #                        ' ,but now {}'.format(self.disk_size_gb))
        # Add a SCSI controller if the one the disk used is not exist.
        if not self.is_there_scsi(scsi_num_tuple[0]):
            self.add_scsi_controller(scsi_num_tuple[0])
        controller_key = 1000 + scsi_num_tuple[0]
        # Make sure the disk is not exist
        for dev in self.vm_obj.config.hardware.device:
            if not isinstance(dev, vim.vm.device.VirtualDisk):
                continue
            if dev.controllerKey == controller_key \
                    and scsi_num_tuple[1] == dev.unitNumber:
                raise RuntimeError('SCSI ({}:{}) already exists'
                                   '.'.format(scsi_num_tuple[0],
                                              scsi_num_tuple[1]))
        new_disk_kb = self.disk_size_gb * 1024 * 1024
        disk_spec = vim.vm.device.VirtualDeviceSpec()
        disk_spec.fileOperation = "create"
        disk_spec.operation = vim.vm.device.VirtualDeviceSpec.Operation.add
        disk_spec.device = vim.vm.device.VirtualDisk()
        disk_spec.device.backing = \
            vim.vm.device.VirtualDisk.FlatVer2BackingInfo()
        # VMWare options
        if self.disk_provision_type == 'thick_lazy':
            pass
        elif self.disk_provision_type == 'thin':
            disk_spec.device.backing.thinProvisioned = True
        elif self.disk_provision_type == 'thick_eager':
            disk_spec.device.backing.eagerlyScrub = True
        else:
            raise RuntimeError('disk_provision_type only can be in '
                               '(thick_lazy, thin, thick_eager).')

        if self.disk_mode not in ('independent_persistent', 'persistent'):
            raise RuntimeError('disk_mode only can be in '
                               '(independent_persistent, persistent).')
        disk_spec.device.backing.diskMode = self.disk_mode
        disk_spec.device.unitNumber = scsi_num_tuple[1]
        disk_spec.device.capacityInKB = new_disk_kb
        disk_spec.device.controllerKey = controller_key

        spec = vim.vm.ConfigSpec()
        spec.deviceChange = [disk_spec]
        task = self.vm_obj.ReconfigVM_Task(spec=spec)
        self._wait_for_tasks([task])
        return task

    def remove_virtual_disk(self, scsi_num_tuple):
        """
         Remove a virtual disk from VM
        Args:
            scsi_num_tuple: scsi number which tuple format.
                            eg. (2,12) means (2:12)

        Returns: task

        """
        assert is_valid_device(scsi_num_tuple)
        virtual_hdd_device = None
        # Find out the disk
        for dev in self.vm_obj.config.hardware.device:
            if (isinstance(dev, vim.vm.device.VirtualDisk)
                    and (1000 + scsi_num_tuple[0]) == dev.controllerKey
                    and scsi_num_tuple[1] == dev.unitNumber):
                virtual_hdd_device = dev
                break
        if not virtual_hdd_device:
            print scsi_num_tuple[0], scsi_num_tuple[1]
            raise RuntimeError('Virtual disk SCSI ({}:{}) could not be found.'
                               .format(scsi_num_tuple[0],
                                       scsi_num_tuple[1]))
        disk_spec = vim.vm.device.VirtualDeviceSpec()
        disk_spec.operation = vim.vm.device.VirtualDeviceSpec.Operation.remove
        disk_spec.device = virtual_hdd_device

        spec = vim.vm.ConfigSpec()
        spec.deviceChange = [disk_spec]
        task = self.vm_obj.ReconfigVM_Task(spec=spec)
        self._wait_for_tasks([task])
        return task

    def add_scsi_controller(self, number):
        """
        Add a SCSI controller to VM
        Args:
            number: The number of a SCSI controller in the range 0~3.

        Returns: task

        """
        controller_key = 1000 + number
        for dev in self.vm_obj.config.hardware.device:
            if isinstance(dev, vim.vm.device.VirtualSCSIController) \
                    and dev.key == controller_key:
                raise RuntimeError('SCSI {} already exists.'
                                   .format(number))
        scsi_cont_spec = vim.vm.device.VirtualDeviceSpec()
        scsi_cont_spec.operation = vim.vm.device.VirtualDeviceSpec.Operation\
                                      .add
        scsi_cont_spec.device = vim.vm.device.VirtualLsiLogicController(
            sharedBus=vim.vm.device.VirtualSCSIController.Sharing.noSharing)
        scsi_cont_spec.device.key = controller_key
        # 3 here is special for vmware
        scsi_cont_spec.device.unitNumber = number + 3
        scsi_cont_spec.device.busNumber = number
        spec = vim.vm.ConfigSpec()
        spec.deviceChange = [scsi_cont_spec]
        task = self.vm_obj.ReconfigVM_Task(spec=spec)
        self._wait_for_tasks([task])
        return task

    def remove_scsi_controller(self, number):
        """
        Remove a SCSI controller from VM
        Args:
            number: The number of a SCSI controller in the range 0~3.

        Returns: task
        """
        if self.is_there_disk(number):
            raise RuntimeError('Could not remove SCSI controller {}.'
                               'Because there is one or more disks under this'
                               ' controller.'.format(number))
        controller_key = 1000 + number
        scsi_ctrl = None
        for dev in self.vm_obj.config.hardware.device:
            if isinstance(dev, vim.vm.device.VirtualSCSIController) \
                    and controller_key == dev.key:
                scsi_ctrl = dev
                break
        if not scsi_ctrl:
            raise RuntimeError('scsi controller {} could not be found'
                               .format(number))
        scsi_cont_spec = vim.vm.device.VirtualDeviceSpec()
        scsi_cont_spec.operation = vim.vm.device.VirtualDeviceSpec.Operation \
                                      .remove
        scsi_cont_spec.device = scsi_ctrl
        spec = vim.vm.ConfigSpec()
        spec.deviceChange = [scsi_cont_spec]
        task = self.vm_obj.ReconfigVM_Task(spec=spec)
        self._wait_for_tasks([task])
        return task

    def is_there_scsi(self, scsi_num):
        """
        Whether there is a certain SCSI controller
        Args:
            scsi_num: The number of a SCSI controller in the range 0~3.

        Returns: SCSI controller device

        """
        for dev in self.vm_obj.config.hardware.device:
            if isinstance(dev, vim.vm.device.VirtualSCSIController):
                if dev.key == 1000 + scsi_num:
                    return dev
        return None

    def is_there_disk(self, scsi_num):
        """
        Whether there are disks under a certain SCSI controller
        Args:
            scsi_num: The number of a SCSI controller in the range 0~3.

        Returns: Disks list

        """
        disk_list = []
        for dev in self.vm_obj.config.hardware.device:
            if isinstance(dev, vim.vm.device.VirtualDisk):
                if dev.controllerKey == 1000 + scsi_num:
                    disk_list.append(dev)
        return disk_list


def is_valid_device(scsi_num_tuple):
    """
    Whether the SCSI number is valid.
    A SCSI controller number is in the range 0~3, and a device unit number
    is in the range 0~15 but except for 7.
    Args:
        scsi_num_tuple: scsi number which tuple format.
                        eg. (2,12) means (2:12)

    Returns: Boolean

    """
    if scsi_num_tuple[1] == 7:
        raise  RuntimeError("Device unit number should not be 7")
    if scsi_num_tuple[0] not in range(MIN_SCSI_NO, MAX_SCSI_NO+1)\
            or scsi_num_tuple[1] not in range(MIN_DEVICE_NO, MAX_DEVICE_NO+1):
        raise RuntimeError("scsi controller number should be from {} to {},"
                           " and device unit number should be from {} to {},"
                           " now it is ({}:{})".format(MIN_SCSI_NO,
                                                       MAX_SCSI_NO,
                                                       MIN_DEVICE_NO,
                                                       MAX_DEVICE_NO,
                                                       scsi_num_tuple[0],
                                                       scsi_num_tuple[1]))
    return True


def is_valid_scsi(scsi_no):
    """
    Whether the SCSI controller number is valid.
    A SCSI controller number is in the range 0~3
    Args:
        scsi_no: scsi controller number

    Returns: Boolean

    """
    if scsi_no not in range(MIN_SCSI_NO, MAX_SCSI_NO+1):
        raise RuntimeError("scsi controller number should be from {} to {},"
                           "but now it is {}".format(MIN_SCSI_NO, MAX_SCSI_NO,
                                                     scsi_no))
    return True


def get_scsi_num_tuple(str_scsi_number):
    """
    Handle the SCSI number string to a tuple format
    Args:
        str_scsi_number:SCSI number string
                        eg. (2:12)

    Returns: a SCSI number tuple

    """
    scsi_num_list = []
    scsi_num_split = str_scsi_number.split(":")
    for scsi_i in scsi_num_split:
        scsi_i = scsi_i.strip('(').strip()
        scsi_i = scsi_i.strip(')').strip()
        scsi_num_list.append(int(scsi_i))
    assert len(scsi_num_list) == 2
    return tuple(scsi_num_list)


def init_virtual_disk(vm_name):
    """
    Create a virtual disk object
    Args:
        vm_name: the name of the virtual machine

    Returns: Virtual disk

    """
    return VirtualDisk(vm_name, disk_provision_type='thin',
                       disk_mode='independent_persistent')


def add_disk_to_vm(virtual_disk, scsi_number,
                   disk_size_gb=MIN_SIZE_STEP_GB,
                   disk_provision_type='thin',
                   disk_mode='independent_persistent'):
    """
    Add disk api for automation cases
    Args:
        virtual_disk: A virtual disk
        scsi_number: SCSI number string
        disk_size_gb: size of the disk to be added
        disk_provision_type: Disk provision type when adding disk for VMWare.
        disk_mode: Disk mode when adding disk for VMWare

    Returns:Boolean

    """
    virtual_disk.disk_size_gb = int(disk_size_gb)
    virtual_disk.disk_provision_type = disk_provision_type
    virtual_disk.disk_mode = disk_mode
    scsi_num_tuple = get_scsi_num_tuple(scsi_number)
    if virtual_disk.add_virtual_disk(scsi_num_tuple):
        return True
    return False


def remove_disk_from_vm(virtual_disk, scsi_number):
    """
    Remove disk api for automation cases
    Args:
        virtual_disk: A virtual disk
        scsi_number: SCSI number string

    Returns:Boolean

    """
    scsi_num_tuple = get_scsi_num_tuple(scsi_number)
    if virtual_disk.remove_virtual_disk(scsi_num_tuple):
        return True
    return False


def remove_scsi_controller_from_vm(virtual_disk, number):
    """
    Remove SCSI controller api for automation cases
    Args:
        virtual_disk: A virtual disk
        number: SCSI controller number

    Returns:Boolean

    """
    number = int(number)
    assert is_valid_scsi(number)
    if virtual_disk.remove_scsi_controller(number):
        return True
    return False


def add_scsi_controller_to_vm(virtual_disk, number):
    """
    Add SCSI controller api for automation cases
    Args:
        virtual_disk: A virtual disk
        number: SCSI controller number

    Returns:Boolean

    """
    number = int(number)
    assert is_valid_scsi(number)
    if virtual_disk.add_scsi_controller(number):
        return True
    return False
